from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Literal
from openai import OpenAI
from fastapi.responses import JSONResponse
import logging
import uvicorn
import httpx
from starlette.background import BackgroundTask
import base64
from config import settings


app = FastAPI()

SYSTEM_PROMPT = '''
<身份>
你是一名擅长中学数学的老师，你需要根据上传的数学题目，提取出题目和问题中包含的所有图形（圆、四边形、三角形等）、点、角、直线（对称轴、切线等）、线段和函数并以规定格式列出。
若题目中元素的信息在问题中需要求解得出（点的坐标、函数的解析式等），你必须计算答案并与其他元素一起列出，以保证输出答案的精确。
</身份>

<需求>
•   题目和各个小问中提到的所有元素包括计算过程中涉及的辅助线都必须包含，以保证输出答案的精确。
•   若题目中需要分类讨论，相关元素的所有情况都必须列出，以保证输出答案的精准。
•   输出严格参照下方格式，不得添加任何额外内容。
</需求>

<输出格式>
### 1. **点**

### 2. **角**

### 3. **直线**

### 4. **线段**

### 5. **函数**

### 6. **图形**
</输出格式>
'''

RECEIVER_URL = "http://0.0.0.0:9000/receive_chunk_stream"

class ProblemInput(BaseModel):
    problem: str
    problem_img: str

class ElementsOutput(BaseModel):
    reasoning: str
    elements: str

logging.basicConfig(
    filename="element_extraction_api.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)


@app.post("/extract_elements", response_model=ElementsOutput)
async def extract_elements(input: ProblemInput):
    try:
        logging.info(f"Received problem: {input.problem}")

        client = OpenAI(api_key=settings.deepseek_api_key, base_url=settings.deepseek_base_url)
        response = client.chat.completions.create(
            model="deepseek-reasoner",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": input.problem},
            ],
            stream=False
        )
        reasoning_content = response.choices[0].message.reasoning_content
        logging.info(f"Reasoning content: {reasoning_content}")
        elements = response.choices[0].message.content
        logging.info(f"Extracted elements: {elements}")
        return JSONResponse(content={"reasoning": reasoning_content, "elements": elements}, status_code=200)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    

def push_stream_to_receiver(problem: str):
    try:
        client = OpenAI(api_key=settings.deepseek_api_key, base_url=settings.deepseek_base_url)
        response = client.chat.completions.create(
            model="deepseek-reasoner",
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": problem},
            ],
            stream=True
        )
        for chunk in response:
            delta = chunk.choices[0].delta
            data = {}
            if delta.reasoning_content:
                logging.info(f"Reasoning content: {delta.reasoning_content}")
                data = {"type": "reasoning", "content": delta.reasoning_content}
            elif delta.content:
                logging.info(f"Extracted elements: {delta.content}")
                data = {"type": "element", "content": delta.content}

            try:
                httpx.post(RECEIVER_URL, json=data, timeout=3.0)
            except Exception as e:
                logging.error(f"Failed to send chunk to receiver: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/extract_elements_stream")
async def extract_elements_proxy(input: ProblemInput):
    logging.info(f"Received problem: {input.problem}")
    task = BackgroundTask(push_stream_to_receiver, problem=input.problem)
    return JSONResponse(content={"status": "accepted", "message": "Streaming started to receiver"}, background=task)


@app.post("/extract_elements_image", response_model=ElementsOutput)
async def extract_elements_image(input: ProblemInput):
    try:
        if input.problem:
            logging.info(f"Received problem: {input.problem}")
        logging.info(f"Received image: {input.problem_img}")

        with open(input.problem_img, "rb") as f:
            base64_str = base64.b64encode(f.read()).decode("ascii")
            base64_url = f"data:image/jpeg;base64,{base64_str}"
            
        client = OpenAI(api_key="4qNfUujEuOkN0WWP9TNHRDfedWM6KWVE8KnErfKRicGBfj9JWkKA4QYktcqLPr4Ch", base_url="https://api.stepfun.com/v1")
        response = client.chat.completions.create(
            model="step-r1-v-mini",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": [
                    {"type": "image_url", "image_url": {"url": base64_url, "detail": "high"}},
                    {"type": "text", "text": input.problem}
                ]},
            ],
            stream=False
        )
        reasoning_content = response.choices[0].message.reasoning
        
        logging.info(f"Reasoning content: {reasoning_content}")
        elements = response.choices[0].message.content
        logging.info(f"Extracted elements: {elements}")
        return JSONResponse(content={"reasoning": reasoning_content, "elements": elements}, status_code=200)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


def push_stream_to_receiver_image(problem: str, problem_img: str):
    try:
        client = OpenAI(api_key="4qNfUujEuOkN0WWP9TNHRDfedWM6KWVE8KnErfKRicGBfj9JWkKA4QYktcqLPr4Ch", base_url="https://api.stepfun.com/v1")
        response = client.chat.completions.create(
            model="step-r1-v-mini",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": [
                    {"type": "image_url", "image_url": {"url": problem_img, "detail": "high"}},
                    {"type": "text", "text": problem}
                ]},
            ],
            stream=True
        )
        for chunk in response:
            delta = chunk.choices[0].delta
            data = {}
            if delta.reasoning:
                logging.info(f"Reasoning content: {delta.reasoning}")
                data = {"type": "reasoning", "content": delta.reasoning}
            elif delta.content:
                logging.info(f"Extracted elements: {delta.content}")
                data = {"type": "element", "content": delta.content}

            try:
                httpx.post(RECEIVER_URL, json=data, timeout=3.0)
            except Exception as e:
                logging.error(f"Failed to send chunk to receiver: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/extract_elements_image_stream")
async def extract_elements_proxy(input: ProblemInput):
    if input.problem:
        logging.info(f"Received problem: {input.problem}")
    logging.info(f"Received image: {input.problem_img}")

    with open(input.problem_img, "rb") as f:
        base64_str = base64.b64encode(f.read()).decode("ascii")
        base64_url = f"data:image/jpeg;base64,{base64_str}"
    
    task = BackgroundTask(push_stream_to_receiver_image, problem=input.problem, problem_image=base64_url)
    return JSONResponse(content={"status": "accepted", "message": "Streaming started to receiver"}, background=task)


if __name__ == "__main__":
    uvicorn.run("element_extraction_api:app", host="0.0.0.0", port=8001, reload=False)